import acl
import sys
import numpy as np
sys.path.append("..")
from CommonModule.ConstModule import PIXEL_FORMAT_YUV_SEMIPLANAR_420, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_420, \
                                     PIXEL_FORMAT_YUV_SEMIPLANAR_422, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_422, \
                                     PIXEL_FORMAT_YUV_SEMIPLANAR_444, \
                                     PIXEL_FORMAT_YVU_SEMIPLANAR_444, \
                                     STATIC_BATCH_RESOLUTION, DYNAMIC_BATCH, \
                                     DYNAMIC_RESOLUTION, \
                                     APPEND, REMOVE, \
                                     DEBUG, INFO, WARNING, ERROR, \
                                     DEVICE, HOST, DVPP,\
                                     ACL_MEM_MALLOC_HUGE_FIRST,\
                                     ACL_MEM_MALLOC_HUGE_ONLY ,\
                                     ACL_MEM_MALLOC_NORMAL_ONLY

from CommonModule.LogModule import DemoLogger, check_ret

class SimpleMemoryPool():
    def __init__(self, mode=STATIC_BATCH_RESOLUTION):
        self.mode = mode
        self.class_name = self.__class__.__name__
        
    def buffer_size_calculation(self, width_stride: int, height_stride: int, data_format: int):
        #print('data_format:', data_format)
        if data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_420 or \
           data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_420:
            buffer_size = width_stride * height_stride * 3 // 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_422 or \
             data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_422:
            buffer_size = width_stride * height_stride * 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_444 or \
             data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_444:
            buffer_size = width_stride * height_stride * 3
        else:
            DemoLogger.do_log(ERROR, self.class_name, 'buffer_size_calculation', \
                    'data_format is not supported')
        return buffer_size

    def stride_calculation(self, height, width, func):
        if func == 'jpegd':
            width_alignment = 128
            height_alignment = 16
        elif func == 'vdec' or \
             func == 'resize' or \
             func == 'crop' or \
             func == 'crop_and_paste':
            width_alignment = 16
            height_alignment = 2
        else:
            DemoLogger.do_log(ERROR, self.class_name ,'stride_calculation', 'func is not supported' )
        width_stride = ((width + width_alignment - 1) // width_alignment) * width_alignment
        height_stride = ((height + height_alignment - 1) // height_alignment) * height_alignment
        return width_stride, height_stride
    
    def memory_allocation(self, func: str, height: int, width: int, data_format: int):
        width_stride, height_stride = self.stride_calculation(height, width, func)
        buffer_size = self.buffer_size_calculation(width_stride, height_stride, data_format)
        start_ptr, ret = acl.media.dvpp_malloc(buffer_size)
        check_ret(self.class_name, 'memory_allocation','fail to malloc dvpp memory',ret)
        return start_ptr, buffer_size, width_stride, height_stride
    
    def memory_allocation_(self, size: int, place: int):
        if place == DVPP:
           ptr, ret = acl.media.dvpp_malloc(size)
           check_ret(self.class_name, 'memory_allocation','fail to malloc dvpp memory',ret)
        elif place == DEVICE:
           ptr, ret = acl.rt.malloc(size, ACL_MEM_MALLOC_HUGE_FIRST)
           check_ret(self.class_name, 'memory_allocation','fail to malloc device memory',ret)
        elif place == HOST:
           ptr, ret = acl.rt.malloc_host(size)
           check_ret(self.class_name, 'memory_allocation','fail to malloc host memory',ret)       
        else:
           DemoLogger.do_log(ERROR, self.class_name, 'memory_allocation', 'out of options')
        return ptr
    def memory_reset(self, ptr, size, value):
        ret = acl.rt.memset(ptr, size, value, size)
        check_ret(self.class_name, 'memory_reset', 'fail to reset by acl.rt.memset', ret)

    def release_memory(self, buffer_ptr, place: int):
        if buffer_ptr != None:
            if place == DVPP:
                ret = acl.media.dvpp_free(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release dvpp memory', ret)
            elif place == DEVICE:
                ret = acl.rt.free(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release memory', ret)
            elif place == HOST:
                ret = acl.rt.free_host(buffer_ptr)
                check_ret(self.class_name, 'release_memory', 'fail to release memory', ret) 
#--------------------------------No used--------------------------------------------------  
class MemoryInfo():
    def __init__(self, mode=STATIC_BATCH_RESOLUTION):
        self.memory_buffer_info = {'buffer_ptr': None, 'total_size': 0, 'occupied_size': 0}
        self.memory_dvpp_info = {}
        self.memory_vpc_info = {}
        self.mode = mode

    def dvpp_memory_info_update(self, func_name, buffer_ptr, buffer_size, batch, flag):
        if self.mode == STATIC_BATCH_RESOLUTION:
            if flag == APPEND:
                self.memory_dvpp_info.update({func_name: \
                    {'buffer_ptr': None, 'buffer_size': [0, 0], 'batch': 0}})
            elif flag == REMOVE:
                del self.memory_dvpp_info[func_name]
            else:
                pass
        else:
            print('Error: Now, the modes beside of STATIC_BATCH_RESOLUTION are not supportted')

    def vpc_memory_info_update(self, func_name, buffer_ptr, buffer_size, batch, in_or_out):
        if self.mode == STATIC_BATCH_RESOLUTION:
            if flag == APPEND:
                self.memory_vpc_info.update({in_or_out: \
                    {'buffer_ptr': None, 'buffer_size': [0, 0], 'batch': 0}})
            elif flag == REMOVE:
                del self.memory_vpc_info[func_name]
            else:
                pass
        else:
            print('Error: Now, the modes beside of STATIC_BATCH_RESOLUTION are not supportted')

class MemoryPool():
    def __init__(self,func_list, mode=STATIC_BATCH_RESOLUTION):
        self.mode = mode
        self.info = {}
    
    def init_memory(self, buffer_size, context):
        acl.rt.set_context(context)
        self.info.update(context,MemoryInfo(mode))
        self.info[context].memory_buffer_info['buffer_ptr'], ret = acl.media.dvpp_malloc(buffer_size)
        self.info[context].memory_buffer_info['total_size'] = (buffer_size//32)*32 + 32

    def buffer_size_calculation(self, width_stride: int, height_stride: int, data_format: int):
        if data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_420 or \
           data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_420:
            buffer_size = width_stride * height_stride * 3 // 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_422 or \
             data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_422:
            buffer_size = width_stride * height_stride * 2
        elif data_format == PIXEL_FORMAT_YUV_SEMIPLANAR_444 or \
             data_format == PIXEL_FORMAT_YVU_SEMIPLANAR_444:
            buffer_size = width_stride * height_stride * 3
        else:
            print('Error: data_format is not supported')
        return buffer_size

    def stride_calculation(self, height, width, func):
        if func == 'jpegd':
            width_alignment = 128
            height_alignment = 16
        elif func == 'vdec':
            width_alignment = 16
            height_alignment = 2
        elif func == 'resize' or \
             func == 'crop' or \
             func == 'crop_and_paste':
            width_alignment = 16
            height_alignment = 2
        else:
            print('Error: wrong func name: ' + func.__name__)
            return -1, -1
        width_stride = ((width + width_alignment - 1) // width_alignment) * width_alignment
        height_stride = ((height + height_alignment - 1) // height_alignment) * height_alignment
        return width_stride, height_stride
    
    def inital_memory_alignment(self, ptr, align=128):
        aligned_ptr = ((ptr + align -1) // align )* align
        return aligned_ptr
        
    def jpeg_decode_memory(self, batch: int, height: int, width: int, data_format: int, \
                           context, need_malloc: bool, start_ptr=None):
        width_stride, height_stride = stride_calculation(height, width, 'jpegd')
        buffer_size = buffer_size_calculation(width_stride, height_stride, data_format)
        if need_malloc:
            start_ptr, ret = acl.media.dvpp.malloc(buffer_size)
            if ret:
                print("Error: fail to malloc acl.media.dvpp memory")
                return ret
        else:
            if start_ptr == None:
                print("Error: no memory location")
                return 
        self.info[context].dvpp_memory_info_update('jpegd', start_ptr, buffer_size, batch, APPEND)

    def video_decode_memory(self, batch: int, height: int, width: int, data_format: int, \
                           context, need_malloc:bool, start_ptr=None):
        width_stride, height_stride = stride_calculation(height, width, 'vdec')
        buffer_size = buffer_size_calculation(width_stride, height_stride, data_format)
        if need_malloc:
            start_ptr, ret = acl.media.dvpp.malloc(buffer_size)
            if ret:
                print("Error: fail to malloc acl.media.dvpp memory")
                return ret
        else:
            if start_ptr == None:
                print("Error: no memory location")
                return 
        self.info[context].dvpp_memory_info_update('vdec', start_ptr, buffer_size, batch, APPEND)

    def resize_memory(self, batch: int, height: int, width: int, data_format: int, \
                      context, need_malloc: bool, start_ptr=None):
        width_stride, height_stride = stride_calculation(height, width, 'resize')
        buffer_size = buffer_size_calculation(width_stride, height_stride, data_format)
        if need_malloc:
            start_ptr, ret = acl.media.dvpp.malloc(buffer_size)
            if ret:
                print("Error: fail to malloc acl.media.dvpp memory")
                return ret
        else:
            if start_ptr == None:
                print("Error: no memory location")
                return 
        self.info[context].dvpp_memory_info_update('resize', start_ptr, buffer_size, batch, APPEND)

    def crop_memory(self, batch: int, height: int, width: int, data_format: int, \
                      context, need_malloc:bool, start_ptr=None):
        width_stride, height_stride = stride_calculation(height, width, 'crop')
        buffer_size = buffer_size_calculation(width_stride, height_stride, data_format)
        if need_malloc:
            start_ptr, ret = acl.media.dvpp.malloc(buffer_size)
            if ret:
                print("Error: fail to malloc acl.media.dvpp memory")
                return ret
        else:
            if start_ptr == None:
                print("Error: no memory location")
                return 
        self.info[context].dvpp_memory_info_update('crop', start_ptr, buffer_size, batch, APPEND)

    def crop_and_paste_memory(self, batch: int, height: int, width: int, data_format: int, \
                      context, need_malloc:bool, start_ptr=None):
        width_stride, height_stride = stride_calculation(height, width, 'crop_and_paste')
        buffer_size = buffer_size_calculation(width_stride, height_stride, data_format)
        if need_malloc:
            start_ptr, ret = acl.media.dvpp.malloc(buffer_size)
            if ret:
                print("Error: fail to malloc acl.media.dvpp memory")
                return ret
        else:
            if start_ptr == None:
                print("Error: no memory location")
                return 
        self.info[context].dvpp_memory_info_update('crop_and_paste', start_ptr, buffer_size, batch, APPEND)

    def device_input_memory(self, index: int, height: int, width: int, data_format: int, need_malloc: bool):
        pass

    def device_output_memory(self, index: int, height: int, width: int, data_format: int, need_malloc: bool):
        pass

    def host_memory(self, index: int, height: int, width: int, data_format: int, need_malloc: bool):
        pass

    def memory_reset(self, ptr, size, value):
        ret = acl.rt.memset(ptr, size, value, size)
        if ret:
            print("Error: fail to reset memory")

    def release_memory(self):
        context, ret = acl.rt.get_context()
        if self.info[context].memory_buffer_info['buffer_ptr'] != None:
            ret = acl.media.dvpp_free(self.info[context].memory_buffer_info['buffer_ptr'])

'''
    def release_memory(self, index: int, func):
        if func == 'jpegd':
            if self.info[context].memory_dvpp_info['jpegd']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.jpeg_memory[index]['output']['buffer_ptr'])
                self.jpeg_memory[index]['output']['buffer_ptr'] = None
                self.jpeg_memory[index]['output']['buffer_size'] = 0
        elif func.__name__ == 'video_decode_memory':
            if self.vdec_memory[index]['output']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.vdec_memory[index]['output']['buffer_ptr'])
                self.vdec_memory[index]['output']['buffer_ptr'] = None
                self.vdec_memory[index]['output']['buffer_size'] = 0
        elif func.__name__ == 'resize_memory':
            if self.resize_memory[index]['output']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.resize_memory[index]['output']['buffer_ptr'])
                self.resize_memory[index]['output']['buffer_ptr'] = None
                self.resize_memory[index]['output']['buffer_size'] = 0
        elif func.__name__ == 'crop_memory':
            if self.crop_memory[index]['output']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.crop_memory[index]['output']['buffer_ptr'])
                self.crop_memory[index]['output']['buffer_ptr'] = None
                self.crop_memory[index]['output']['buffer_size'] = 0
        elif func.__name__ == 'crop_and_paste_memory':
            if self.crop_and_paste_memory[index]['output']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.crop_and_paste_memory[index]['output']['buffer_ptr'])
                self.crop_and_paste_memory[index]['output']['buffer_ptr'] = None
                self.crop_and_paste_memory[index]['output']['buffer_size'] = 0
        elif func.__name__ == 'device_input_memory':
            if self.device_input_memory[index]['output']['buffer_ptr'] != None:
                ret = acl.media.dvpp_free(self.device_memory[index]['output']['buffer_ptr'])
                self.device_input_memory[index]['output']['buffer_ptr'] = None
                self.device_input_memory[index]['output']['buffer_size'] = 0
        print(ret)
'''
#-----------------------------------------------------------------------------------------------------





